/* The LibVMI Library is an introspection library that simplifies access to
 * memory in a target virtual machine or in a file containing a dump of
 * a system's physical memory.  LibVMI is based on the XenAccess Library.
 *
 * Copyright 2011 Sandia Corporation. Under the terms of Contract
 * DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government
 * retains certain rights in this software.
 *
 * Author: Bryan D. Payne (bdpayne@acm.org)
 *
 * This file is part of LibVMI.
 *
 * LibVMI is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Lesser General Public License as published by the
 * Free Software Foundation, either version 3 of the License, or (at your
 * option) any later version.
 *
 * LibVMI is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
 * License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with LibVMI.  If not, see <http://www.gnu.org/licenses/>.
 */
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/mman.h>
#include <stdio.h>
#include <inttypes.h>
#include <wchar.h>
#include <libvmi/libvmi.h>
unsigned long DWORD_ALIGN( offset, base ) 
{
	return( ( (offset + base + 3) & 0xfffffffcL) - (base & 0xfffffffcL) );
}
unicode_string_t * direct_read_widechar( vmi_instance_t vmi, addr_t vaddr, vmi_pid_t pid ) 
{
	access_context_t _ctx = {
			.translate_mechanism	= VMI_TM_PROCESS_PID,
			.addr			= vaddr,
			.pid			= pid
		};
	size_t			buffer_len	= 0;
	unicode_string_t	*us		= 0;
	/* return val */
	uint16_t		value;
	vmi_read_16_va( vmi, vaddr, pid, &value );
	while ( value != 0x0000 ) 
	{
		buffer_len += 0x02;
		vmi_read_16_va( vmi, vaddr + buffer_len, pid, &value );
	}
	/* printf("buffer-->%d,value-->%d\n",buffer_len,value); */
	us		= malloc( sizeof(unicode_string_t) );
	us->length	= buffer_len;
	us->contents	= malloc( sizeof(uint8_t) * (buffer_len + 2) );
	if ( !us->contents )
			goto out_error;
	if ( VMI_FAILURE == vmi_read( vmi, &_ctx, us->length, us->contents, NULL ) ) 
	{
		/* dbprint(VMI_DEBUG_READ, "--%s: failed to read UNICODE_STRING buffer\n",__FUNCTION__); */
		printf( "read total string failed\n" );
		goto out_error;
	}
	/* if */
	/* end with NULL (needed?) */
	us->contents[buffer_len]	= 0;
	us->contents[buffer_len + 1]	= 0;
	us->encoding = "UTF-16";
	return(us);
	out_error:
		if ( us )
			free( us->contents );
	free( us );
	return(0);
}
int main( int argc, char **argv ) 
{
	vmi_instance_t	vmi;
	addr_t		list_head	= 0, cur_list_entry = 0, next_list_entry = 0;
	addr_t		current_process = 0;
	char		*procname	= NULL;
	vmi_pid_t	pid		= 0;
	unsigned long	tasks_offset	= 0, pid_offset = 0, name_offset = 0;
	unsigned long	win_peb		= 0x1a8, win_ldr = 0x00c, win_module = 0x00c;
	unsigned long	win_module2	= 0x014, win_module3 = 0x01c;
	status_t	status;
	/* this is the VM or file that we are looking at */
	if ( argc != 2 ) 
	{
		printf( "Usage: %s <vmname>\n", argv[0] );
		return(1);
	}
	/* if */
	char *name = argv[1];
	/* initialize the libvmi library */
	if ( VMI_FAILURE ==
		     vmi_init_complete( &vmi, name, VMI_INIT_DOMAINNAME, NULL,
					VMI_CONFIG_GLOBAL_FILE_ENTRY, NULL, NULL ) ) 
	{
		printf( "Failed to init LibVMI library.\n" );
		return(1);
	}
	/* init the offset values */
	if ( VMI_OS_LINUX == vmi_get_ostype( vmi ) ) 
	{
		if ( VMI_FAILURE == vmi_get_offset( vmi, "linux_tasks", &tasks_offset ) )
					goto error_exit;
		if ( VMI_FAILURE == vmi_get_offset( vmi, "linux_name", &name_offset ) )
					goto error_exit;
		if ( VMI_FAILURE == vmi_get_offset( vmi, "linux_pid", &pid_offset ) )
					goto error_exit;
	} else if ( VMI_OS_WINDOWS == vmi_get_ostype( vmi ) ) 
	{
		if ( VMI_FAILURE == vmi_get_offset( vmi, "win_tasks", &tasks_offset ) )
					goto error_exit;
		if ( VMI_FAILURE == vmi_get_offset( vmi, "win_pname", &name_offset ) )
					goto error_exit;
		if ( VMI_FAILURE == vmi_get_offset( vmi, "win_pid", &pid_offset ) )
					goto error_exit;
	}
	/* pause the vm for consistent memory access */
	if ( vmi_pause_vm( vmi ) != VMI_SUCCESS ) 
	{
		printf( "Failed to pause VM\n" );
		goto error_exit;
	}
	/* if */
	/* demonstrate name and id accessors */
	char		*name2 = vmi_get_name( vmi );
	vmi_mode_t	mode;
	if ( VMI_FAILURE == vmi_get_access_mode( vmi, NULL, 0, NULL, &mode ) )
			goto error_exit;
	if ( VMI_FILE != mode ) 
	{
		uint64_t id = vmi_get_vmid( vmi );
		printf( "Process listing for VM %s (id=%" PRIu64 ")\n", name2, id );
	} else 
	{
		printf( "Process listing for file %s\n", name2 );
	}
	free( name2 );
	os_t os = vmi_get_ostype( vmi );
	/* get the head of the list */
	if ( VMI_OS_LINUX == os ) 
	{
		/* Begin at PID 0, the 'swapper' task. It's not typically shown by OS
		 *  utilities, but it is indeed part of the task list and useful to
		 *  display as such.
		 */
		if ( VMI_FAILURE == vmi_translate_ksym2v( vmi, "init_task", &list_head ) )
					goto error_exit;
		list_head += tasks_offset;
	} else if ( VMI_OS_WINDOWS == os ) 
	{
		/* find PEPROCESS PsInitialSystemProcess */
		if ( VMI_FAILURE == vmi_read_addr_ksym( vmi, "PsActiveProcessHead", &list_head ) ) 
		{
			printf( "Failed to find PsActiveProcessHead\n" );
			goto error_exit;
		}
	}
	cur_list_entry = list_head;
	if ( VMI_FAILURE == vmi_read_addr_va( vmi, cur_list_entry, 0, &next_list_entry ) ) 
	{
		printf( "Failed to read next pointer in loop at %" PRIx64 "\n", cur_list_entry );
		goto error_exit;
	}
	/* walk the task list */
	while ( 1 ) 
	{
		current_process = cur_list_entry - tasks_offset;
		/* Note: the task_struct that we are looking at has a lot of
		 * information.  However, the process name and id are burried
		 * nice and deep.  Instead of doing something sane like mapping
		 * this data to a task_struct, I'm just jumping to the location
		 * with the info that I want.  This helps to make the example
		 * code cleaner, if not more fragile.  In a real app, you'd
		 * want to do this a little more robust :-)  See
		 * include/linux/sched.h for mode details */
		/* NOTE: _EPROCESS.UniqueProcessId is a really VOID*, but is never > 32 bits,
		 * so this is safe enough for x64 Windows for example purposes */
		vmi_read_32_va( vmi, current_process + pid_offset, 0, (uint32_t *) &pid );
		procname = vmi_read_str_va( vmi, current_process + name_offset, 0 );
		unsigned long temp = 0;
		vmi_read_addr_va( vmi, current_process + win_peb, 0, &temp );
		/*
		 * vmi_translate_uv2p(vmi, temp, 2888,temp2);
		 * temp = temp2;
		 */
		vmi_read_addr_va( vmi, temp + win_ldr, pid, &temp );
		vmi_read_addr_va( vmi, temp + win_module, pid, &temp );
		if ( strcmp( procname, "notepad.exe" ) == 0 ) 
		{
			printf( "dll library:\n" );
			/* printf("%"PRIx64"\n",temp); */
			unicode_string_t	*us = NULL;
			unicode_string_t	out;
			us = vmi_read_unicode_str_va( vmi, temp + 0x024, pid );
			vmi_convert_str_encoding( us, &out, "UTF-8" );
			printf( "%s\n", out.contents );
			while ( strcmp( out.contents, "" ) != 0 ) 
			{
				printf("------------------------------------------------------------------------\n");
				vmi_read_addr_va( vmi, temp, pid, &temp );
				/* printf("%"PRIx64"\n",temp); */
				us = vmi_read_unicode_str_va( vmi, temp + 0x024, pid );
				vmi_convert_str_encoding( us, &out, "UTF-8" );
				printf( "%s\n", out.contents );
				/*
				 * printf("temp--->%"PRIx64"\n",temp+0x24);
				 * return 0;
				 */
				if(strcmp(out.contents,"C:\\Windows\\system32\\SHELL32.dll") != 0)
								//if ( 1 ) 
				{
					/* get the dllBase address "dllBase" */
					unsigned long dllBase = 0, e_lfanew = 0, rsrc_virtAddress = 0;
					vmi_read_addr_va( vmi, temp + 0x018, pid, &dllBase );
					vmi_read_addr_va( vmi, dllBase + 0x03c, pid, &e_lfanew );
					//printf( "%" PRIx64 "\n", dllBase );
					//printf("%s\n",out.contents);
					/* MayBe modifying later */
					vmi_read_addr_va( vmi, dllBase + e_lfanew + 0x018 + 0x060 + 0x008 + 0x008, pid, &rsrc_virtAddress );
					//printf( "%" PRIx64 "\n", rsrc_virtAddress );
					/*
					 * get NumberOfNamedEntries and NumberOfIdEntries: dt _IMAGE_RESOURCE_DIRECTORY dllBase+resr_virtASddress
					 * 0x00c-->NamedEntries:integer(may be)
					 * 0x00e-->IdEntries:integer(may be)
					 */
					//return 0;
					/*
					 * IMAGE_RESOURCE_DIRECTORY's root size 0x010, IMAGE_RESOURCE_DIRECTORY_ENTRY size 0x008
					 * rsrc_virtAddress = 0x240;
					 */
					uint16_t numberOfNamedEntries = 0;
					uint16_t numberOfIdEntries = 0;
					vmi_read_16_va( vmi, dllBase + rsrc_virtAddress + 0x00c, pid, &numberOfNamedEntries );
					vmi_read_16_va( vmi, dllBase + rsrc_virtAddress + 0x00e, pid, &numberOfIdEntries );
					printf( "numberOfNamedEntris--->%d,numberOfIdEntris--->%d\n", numberOfNamedEntries, numberOfIdEntries );
					/* unsigned long tempEntry = dllBase+rsrc_virtAddress+0x010+0x008+0x008; */
					unsigned long	tempEntry		= dllBase + rsrc_virtAddress + 0x010;
					unsigned long	entryFirstMember	= 0x001, entrySecondMember = 0;
					/* int aaaa=0; */
					//printf( "*****STARTS*****\n" );
					/*
 * size_t outNum = 0;
 * unsigned long aa = 0;
 * status_t testStatus =  vmi_read_va(vmi,tempEntry,pid,4,&aa,&outNum);
 * printf("aaa%"PRIx64"\n",aa);
 */
					/*
					 * unsigned long index = 0;
					 * for(index = 0x7797DF00;index<0x7797E0F0;index+=9)//0x7797E0F0 0x779DBFF0
					 * {
					 *  status_t result = vmi_read_addr_va(vmi,index,pid,&entryFirstMember);
					 * if(result == 0){
					 *  printf("index--->%"PRIx64",result--->%d\n",index,result);
					 * }
					 */
					/* } */
					status_t result = vmi_read_addr_va( vmi, tempEntry, pid, &entryFirstMember );
					if ( VMI_FAILURE == result ) 
					{
						printf( "Probably current dll not in the memory ~_~\n" );
						continue;
					}
					//printf( "result--->%d\n", result );
					//printf( "%" PRIx64 "\n", tempEntry );
					//printf( "%" PRIx64 "\n", entryFirstMember );
					/* dllBase = 76e10000;rsrv_virt = e0000; */
					int i = 1;
					while ( entryFirstMember != 0x10 && i <= numberOfNamedEntries + numberOfIdEntries ) 
					{
						i++;
						tempEntry += 0x008;
						vmi_read_addr_va( vmi, tempEntry, pid, &entryFirstMember );
					}
					if ( entryFirstMember == 0x10 ) 
					{
						/* printf("find it-=---->\n"); */
						vmi_read_addr_va( vmi, tempEntry + 0x004, pid, &entrySecondMember );
					} else 
					{
						printf("Can not find rsrc root node");
						continue;
					}
					/* vmi_read_addr_va(vmi,tempEntry+0x004,pid,&entrySecondMember); */
					//printf( "%" PRIx64 "\n", entrySecondMember );
					unsigned long tempOffsetToDirectory = 0;
					tempOffsetToDirectory = (entrySecondMember & 0x7fffffff);
					/* 0x58 */
					//printf( "%" PRIx64 "\n", tempOffsetToDirectory );
					vmi_read_16_va( vmi, dllBase + rsrc_virtAddress + tempOffsetToDirectory + 0x00c, pid, &numberOfNamedEntries );
					vmi_read_16_va( vmi, dllBase + rsrc_virtAddress + tempOffsetToDirectory + 0x00e, pid, &numberOfIdEntries );
					printf( "totalCount-->%d\n", numberOfNamedEntries + numberOfIdEntries );
					/* step entries num */
					vmi_read_addr_va( vmi, dllBase + rsrc_virtAddress + tempOffsetToDirectory + 0x010 + 0x004, pid, &entrySecondMember );
					/*
					 * printf("%"PRIx64"\n",entrySecondMember);
					 * tempOffsetToDirectory = (entrySecondMember & 0x7fffffff);//0xa0
					 * printf("%"PRIx64"\n",tempOffsetToDirectory);
					 * vmi_read_16_va(vmi,dllBase+rsrc_virtAddress+tempOffsetToDirectory+0x00c,pid,&numberOfNamedEntries);
					 * vmi_read_16_va(vmi,dllBase+rsrc_virtAddress+tempOffsetToDirectory+0x00e,pid,&numberOfIdEntries);
					 * printf("totalCount-->%d\n",numberOfNamedEntries+numberOfIdEntries);
					 */
					/*
					 * get NumberOfNamedEntries and NumberOfIdEntries: dt _IMAGE_RESOURCE_DIRECTORY dllBase+resr_virtASddress+tempOffsetToDirectory
					 * 0x00c-->NamedEntries:integer(may be)
					 * 0x00e-->IdEntries:integer(may be)
					 */
					tempEntry = dllBase + rsrc_virtAddress + tempOffsetToDirectory + 0x010;
					vmi_read_addr_va( vmi, tempEntry, pid, &entryFirstMember );
					vmi_read_addr_va( vmi, tempEntry + 0x004, pid, &entrySecondMember );
					tempOffsetToDirectory = (entrySecondMember & 0x7fffffff);
					/* 0xa0 */
					//printf( "tempOffsetToDirectory--->%" PRIx64 "\n", tempOffsetToDirectory );
					/*
					 * get NumberOfNamedEntries and NumberOfIdEntries: dt _IMAGE_RESOURCE_DIRECTORY dllBase+resr_virtASddress+tempOffsetToDirectory
					 * 0x00c-->NamedEntries:integer(may be)
					 * 0x00e-->IdEntries:integer(may be)
					 */
					tempEntry = dllBase + rsrc_virtAddress + tempOffsetToDirectory + 0x010;
					vmi_read_addr_va( vmi, tempEntry, pid, &entryFirstMember );
					vmi_read_addr_va( vmi, tempEntry + 0x004, pid, &entrySecondMember );
					tempOffsetToDirectory = (entrySecondMember & 0x7fffffff);
					/* 0xd8 */
					//printf( "tempOffsetToDirectory--->%" PRIx64 "\n", tempOffsetToDirectory );
					/* arrive to the direct pointer */
					tempEntry = dllBase + rsrc_virtAddress + tempOffsetToDirectory;
					vmi_read_addr_va( vmi, tempEntry, pid, &entryFirstMember );
					/* entryFirstMember = 0xe00f0 */
					vmi_read_addr_va( vmi, tempEntry + 0x004, pid, &entrySecondMember );
					tempOffsetToDirectory = (entrySecondMember & 0x7fffffff);
					/* 0x380 */
					unsigned long	data_size	= tempOffsetToDirectory;
					unsigned long	dataStart	= entryFirstMember;
					//printf( "tempOffsetToDirectory--->%" PRIx64 "\n", tempOffsetToDirectory );
					//printf( "entryFirstMember--->%" PRIx64 "\n", entryFirstMember );
					/* return 0; */
					/* get tag_VS_VERSION-->wValueLength */
					unsigned long wValueLength;
					vmi_read_addr_va( vmi, dllBase + entryFirstMember + 0x02, pid, &wValueLength );
					unsigned long childNodeStart = dllBase + entryFirstMember + wValueLength + 0x28;
					//printf( "childNodeStart--->%" PRIx64 "\n", childNodeStart + 0x06 );
					/* vmi_read_addr_va(vmi,childNodeStart+0x06,pid,&wValueLength); */
					unicode_string_t stringFileInfoUTF8;
					/*
					 * unicode_string_t stringFileInfo;
					 * unicode_string_t * stringFileInfo =(unicode_string_t *) vmi_read_str_va(vmi, childNodeStart+0x06, pid);//0x06 step StringFileInfo-->wLength 2;wValueLength 2; wType 2;
					 */
					unicode_string_t * stringFileInfo = direct_read_widechar( vmi, childNodeStart + 0x06, pid );
					vmi_convert_str_encoding( stringFileInfo, &stringFileInfoUTF8, "UTF-8" );
					if ( strcmp( "StringFileInfo", stringFileInfoUTF8.contents ) == 0 ) 
					/* stringFileIno 0x1e */ 
					{
						/* step StringTable-->wLength 2 wVaueLength 2 wType 2 szKey 0x12 */
						uint16_t currentStringLength;
						tempOffsetToDirectory = childNodeStart + 0x06 + 0x06 + 0x12 + 0x1e;
						vmi_read_16_va( vmi, tempOffsetToDirectory, pid, &currentStringLength );
						//printf( "currentStringLength--->%d\n", currentStringLength );
						/* step String-->wLength 2 wValueLength 2 wType 2 */
						tempOffsetToDirectory	+= 0x06;
						stringFileInfo		= direct_read_widechar( vmi, tempOffsetToDirectory, pid );
						vmi_convert_str_encoding( stringFileInfo, &stringFileInfoUTF8, "UTF-8" );
						tempOffsetToDirectory -= 0x06;
						while ( strcmp( "FileVersion", stringFileInfoUTF8.contents ) != 0 && tempOffsetToDirectory - dataStart - dllBase < data_size ) 
						{
							/* tempOffsetToDirectory += stringFileInfo->length + 0x02 + 0x02; */
							/*
							 * stringFileInfo = direct_read_widechar(vmi, tempOffsetToDirectory, pid);
							 *        vmi_convert_str_encoding(stringFileInfo,&stringFileInfoUTF8,"UTF-8");
							 */
							/* tempOffsetToDirectory = childNodeStart+0x06+0x06+0x12+0x1e; */
							tempOffsetToDirectory += currentStringLength;
							/* tempOffsetToDirectory += 0x06; */
							tempOffsetToDirectory	= DWORD_ALIGN( tempOffsetToDirectory - dllBase, tempOffsetToDirectory - dllBase );
							tempOffsetToDirectory	+= dllBase;
							status_t result1 = vmi_read_16_va( vmi, tempOffsetToDirectory, pid, &currentStringLength );
							if ( result1 != 0 ) 
							{
								return(0);
							}
							tempOffsetToDirectory	+= 0x06;
							stringFileInfo		= direct_read_widechar( vmi, tempOffsetToDirectory, pid );
							vmi_convert_str_encoding( stringFileInfo, &stringFileInfoUTF8, "UTF-8" );
							//printf( "%s\n", stringFileInfoUTF8.contents );
							tempOffsetToDirectory -= 0x06;
							//printf( "tempOffsetToDirectory--->%" PRIx64 "\n", tempOffsetToDirectory );
						}
						if ( tempOffsetToDirectory - dataStart - dllBase >= data_size ) 
						{
							break;
						}
						unsigned long baseString = tempOffsetToDirectory;
						tempOffsetToDirectory	+= 0x06;
						tempOffsetToDirectory	= tempOffsetToDirectory + stringFileInfo->length + 0x02;
						tempOffsetToDirectory	= DWORD_ALIGN( tempOffsetToDirectory - dllBase, baseString - dllBase );
						tempOffsetToDirectory	+= dllBase;
						stringFileInfo		= direct_read_widechar( vmi, tempOffsetToDirectory, pid );
						vmi_convert_str_encoding( stringFileInfo, &stringFileInfoUTF8, "UTF-8" );
					}
					printf( "%s\n", stringFileInfoUTF8.contents );
				}
			}
			printf( "\n" );
			return 0;
		}
		if ( !procname ) 
		{
			printf( "Failed to find procname\n" );
			goto error_exit;
		}
		/* print out the process name */
		if ( strcmp( procname, "notepad.exe" ) == 0 ) 
		{
			printf( "[%5d] %s (struct addr:%" PRIx64 ")\n", pid, procname, current_process );
		}
		if ( procname ) 
		{
			free( procname );
			procname = NULL;
		}
		/* follow the next pointer */
		cur_list_entry	= next_list_entry;
		status		= vmi_read_addr_va( vmi, cur_list_entry, 0, &next_list_entry );
		if ( status == VMI_FAILURE ) 
		{
			printf( "Failed to read next pointer in loop at %" PRIx64 "\n", cur_list_entry );
			goto error_exit;
		}
		/* In Windows, the next pointer points to the head of list, this pointer is actually the
		 * address of PsActiveProcessHead symbol, not the address of an ActiveProcessLink in
		 * EPROCESS struct.
		 * It means in Windows, we should stop the loop at the last element in the list, while
		 * in Linux, we should stop the loop when coming back to the first element of the loop
		 */
		if ( VMI_OS_WINDOWS == os && next_list_entry == list_head ) 
		{
			break;
		} else if ( VMI_OS_LINUX == os && cur_list_entry == list_head ) 
		{
			break;
		}
	}
	;
	error_exit:
	/* resume the vm */
	vmi_resume_vm( vmi );
	/* cleanup any memory associated with the LibVMI instance */
	vmi_destroy( vmi );
	/* printf("win_tasks:%"PRIx64" win_pname:%"PRIx64" win_pid:%"PRIx64"\n", tasks_offset, pid_offset, name_offset); */
	return(0);
}